package com.feidee.fd.sml.algorithm.component.feature

import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.feature.VectorSlicer
import org.apache.spark.sql.DataFrame

/**
  * @Author: dongguosheng
  * @Date: 2019/3/20 19:44
  * @Review songhaicheng
  * @Email: guosheng_dong@sui.com
  */
case class VectorSliceEncoderParam(
                                    override val input_pt: String,
                                    override val output_pt: String,
                                    override val hive_table: String,
                                    override val flow_time: String,
                                    override val inputCol: String,
                                    override val outputCol: String,
                                    override val preserveCols: String,
                                    override val modelPath: String,
                                    // 起始切分索引值，与参数 names 不冲突
                                    indices: Array[Int],
                                    //（不推荐使用）向量中特征的名字，这要求向量列有 AttributeGroup，因为这根据 Attribute 来匹配名字字段
                                    names: Array[String]
                                  ) extends FeatureParam {

  def this() = this(null, null, null, null, "input", "features", null, null, Array.empty[Int], Array.empty[String])

  override def verify(): Unit = {
    super.verify()
    require(indices.length > 0 || names.length > 0, "param names.size or indices.size must be grater than 0")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("names" -> names)
    map += ("indices" -> indices)
    map
  }
}


class VectorSliceEncoder extends AbstractFeatureEncoder[VectorSliceEncoderParam] {
  override def setUp(param: VectorSliceEncoderParam, data: DataFrame): Array[PipelineStage] = {
    val vectorSlicer = new VectorSlicer()
      .setInputCol(param.inputCol)
      .setOutputCol(param.outputCol)

    if (param.indices != null && param.indices.length > 0) {
      vectorSlicer.setIndices(param.indices)
    }
    if (param.names != null && param.names.length > 0) {
      vectorSlicer.setNames(param.names)
    }

    Array(vectorSlicer)
  }
}

object VectorSliceEncoder {

  def apply(paramStr: String): Unit = {
    new VectorSliceEncoder()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    VectorSliceEncoder(args(0))
  }
}